#include "main.h"
#include "engine.h"

int main() {
    // 1. 加载输入图像
    const std::string image_path = "test.jpg";
    cv::Mat input_image = cv::imread(image_path);
    
    if(input_image.empty()) {
        std::cerr << "Error: Could not load image at " << image_path << std::endl;
        return -1;
    }
    
    // 2. 初始化TensorRT引擎
    const std::string engine_path = "/home/robocon/Foursteps/deploy/models/yolov8n.trt";
    try {
        TRTEngine engine(engine_path);
        
        // 3. 准备输出容器（根据模型输出维度）
        std::vector<float> output(8400 * 6); // 8400个检测框 × 6个参数(x,y,w,h,conf,cls)

        // 4. 执行推理
        engine.infer(input_image, output);
        
        // 5. 解析输出结果（需根据模型结构实现）
        /* 示例输出解析伪代码：
        for(int i = 0; i < 8400; ++i) {
            float conf = output[i*6 + 4];
            if(conf > 0.5) { // 置信度阈值
                int cls = output[i*6 + 5];
                float x = output[i*6 + 0];
                float y = output[i*6 + 1];
                float w = output[i*6 + 2];
                float h = output[i*6 + 3];
                // 绘制检测框...
            }
        }
        */
        std::cout << "Inference completed. Output size: " << output.size() << std::endl;
        
    } catch(const std::exception& e) {
        std::cerr << "Runtime error: " << e.what() << std::endl;
        return -2;
    }

    return 0;
}